import asyncio
import json
import aiohttp


class StreamableHTTPMCPClient:
    """Streamable HTTP MCP 客户端实现"""

    def __init__(self, server_url: str = "http://127.0.0.1:8000"):
        self.server_url = server_url
        self.session_id = None
        self.sse_task = None
        self.session = None
        self.server_mode = None

    async def __aenter__(self):
        """异步上下文管理器入口"""
        self.session = aiohttp.ClientSession()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """异步上下文管理器出口"""
        if self.sse_task:
            self.sse_task.cancel()
        if self.session:
            await self.session.close()

    async def initialize(self) -> bool:
        """初始化 MCP 连接"""
        try:
            payload = {
                "jsonrpc": "2.0",
                "id": 1,
                "method": "initialize",
                "params": {
                    "protocolVersion": "2024-11-05",
                    "capabilities": {"tools": {}, "resources": {}, "prompts": {}},
                    "clientInfo": {
                        "name": "Streamable HTTP MCP Client",
                        "version": "1.0.0",
                    },
                },
            }

            async with self.session.post(
                f"{self.server_url}/message", json=payload
            ) as response:
                if response.status == 200:
                    result = await response.json()
                    self.server_mode = result["result"]["serverInfo"]["mode"]

                    if self.server_mode == "stateful":
                        # 有状态模式：从服务器信息中获取会话 ID
                        self.session_id = result["result"]["serverInfo"].get(
                            "session_id"
                        )

                    print(" 连接初始化成功")
                    print(f"   协议版本: {result['result']['protocolVersion']}")
                    print(f"   服务器: {result['result']['serverInfo']['name']}")
                    print(f"   服务器模式: {self.server_mode}")
                    if self.session_id:
                        print(f"   会话 ID: {self.session_id}")
                    return True
                else:
                    print(f"❌ 连接初始化失败: {response.status}")
                    return False

        except Exception as e:
            print(f"❌ 连接初始化异常: {e}")
            return False

    async def list_tools(self) -> list:
        """获取可用工具列表"""
        try:
            payload = {"jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {}}

            # 使用统一的 /message 端点
            url = f"{self.server_url}/message"
            headers = {}
            if self.session_id:
                headers["X-Session-ID"] = self.session_id

            async with self.session.post(
                url, json=payload, headers=headers
            ) as response:
                if response.status == 200:
                    result = await response.json()
                    tools = result["result"]["tools"]
                    print(f"获取到 {len(tools)} 个工具:")
                    for tool in tools:
                        print(f"   - {tool['name']}: {tool['description']}")
                    return tools
                else:
                    print(f"❌ 获取工具列表失败: {response.status}")
                    return []

        except Exception as e:
            print(f"❌ 获取工具列表异常: {e}")
            return []

    async def call_tool(self, tool_name: str, arguments):
        """调用工具"""
        try:
            payload = {
                "jsonrpc": "2.0",
                "id": 3,
                "method": "tools/call",
                "params": {"name": tool_name, "arguments": arguments},
            }

            # 使用统一的 /message 端点
            url = f"{self.server_url}/message"
            headers = {}
            if self.session_id:
                headers["X-Session-ID"] = self.session_id

            async with self.session.post(
                url, json=payload, headers=headers
            ) as response:
                if response.status == 200:
                    content_type = response.headers.get("content-type", "")

                    if "text/event-stream" in content_type:
                        # SSE 流式响应
                        print(f"📡 工具 {tool_name} 返回流式响应...")
                        return await self._handle_sse_response(response)
                    else:
                        # 普通 JSON 响应
                        try:
                            result = await response.json()

                            # 检查是否有错误
                            if "error" in result:
                                print(
                                    f"⚠️  工具 {tool_name} 调用警告: {result['error']['message']}"
                                )
                            else:
                                print(f" 工具 {tool_name} 调用成功")

                            return result
                        except Exception as json_error:
                            print(f"⚠️  JSON 解析失败: {json_error}")
                            # 尝试读取原始响应内容
                            text_content = await response.text()
                            print(f"   原始响应: {text_content[:200]}...")
                            return {
                                "error": "JSON parsing failed",
                                "raw_content": text_content,
                            }
                else:
                    print(f"❌ 工具调用失败: {response.status}")
                    return None

        except Exception as e:
            print(f"❌ 工具调用异常: {e}")
            return None

    async def _handle_sse_response(self, response):
        """处理 SSE 响应"""
        print("📡 接收到 SSE 流式响应...")
        result = None

        try:
            async for line in response.content:
                line = line.decode("utf-8").strip()
                if line.startswith("data: "):
                    try:
                        data = json.loads(line[6:])  # 移除 'data: ' 前缀

                        if "method" in data:
                            # 通知消息
                            if "notifications/progress" in data["method"]:
                                progress_data = data["params"]
                                if progress_data["type"] == "begin":
                                    print(f" 任务开始: {progress_data['title']}")
                                elif progress_data["type"] == "update":
                                    print(f" 进度更新: {progress_data['message']}")
                                elif progress_data["type"] == "end":
                                    print(f" 任务完成: {progress_data['title']}")
                            elif "notifications/heartbeat" in data["method"]:
                                print(
                                    f" 心跳: {data['params'].get('timestamp', 'N/A')}"
                                )
                            elif "notifications/connection" in data["method"]:
                                print(f" 连接状态: {data['params']['status']}")

                        elif "result" in data:
                            # 结果消息
                            result = data
                            print(
                                f"📋 最终结果: {data['result']['content'][0]['text']}"
                            )
                            break

                    except json.JSONDecodeError as e:
                        print(f"⚠️  SSE 数据解析失败: {e}, 原始数据: {line}")
                        continue

        except Exception as e:
            print(f"⚠️  SSE 流处理异常: {e}")

        return result or {"error": "No result received from SSE stream"}

    async def establish_sse_connection(self) -> bool:
        """建立 SSE 连接（仅在有状态模式下可用）"""
        if self.server_mode == "stateless":
            print(" SSE 连接在无状态模式下不可用")
            return False

        if not self.session_id:
            print(" 请先初始化连接")
            return False

        try:
            # 使用统一的 /message 端点
            url = f"{self.server_url}/message"
            headers = {}
            if self.session_id:
                headers["X-Session-ID"] = self.session_id

            async with self.session.get(url, headers=headers) as response:
                if response.status == 200:
                    print(" SSE 连接建立成功")

                    # 启动 SSE 监听任务
                    self.sse_task = asyncio.create_task(self._listen_sse(response))
                    return True
                else:
                    print(f" SSE 连接建立失败: {response.status}")
                    return False

        except Exception as e:
            print(f" SSE 连接异常: {e}")
            return False

    async def _listen_sse(self, response):
        """监听 SSE 消息"""
        try:
            async for line in response.content:
                line = line.decode("utf-8").strip()
                if line.startswith("data: "):
                    try:
                        data = json.loads(line[6:])
                        if "method" in data:
                            print(f"📨 SSE 消息: {data['method']}")
                    except json.JSONDecodeError:
                        continue
        except asyncio.CancelledError:
            print("📡 SSE 连接已取消")
        except Exception as e:
            # 检查是否是正常的连接关闭
            error_msg = str(e).lower()
            if any(
                keyword in error_msg
                for keyword in [
                    "connection closed",
                    "connection reset",
                    "broken pipe",
                    "end of stream",
                    "stream closed",
                ]
            ):
                print("📡 SSE 连接正常关闭")
            else:
                print(f"❌ SSE 监听异常: {e}")

    async def shutdown(self) -> bool:
        """关闭连接"""
        if self.server_mode == "stateless":
            print(" 无状态服务器无需关闭连接")
            return True

        if not self.session_id:
            return True

        try:
            payload = {"jsonrpc": "2.0", "id": 4, "method": "shutdown", "params": {}}

            headers = {"X-Session-ID": self.session_id}
            async with self.session.post(
                f"{self.server_url}/message", json=payload, headers=headers
            ) as response:
                if response.status == 200:
                    print("连接关闭成功")
                    self.session_id = None
                    return True
                else:
                    print(f" 连接关闭失败: {response.status}")
                    return False

        except Exception as e:
            print(f" 连接关闭异常: {e}")
            return False


async def demo_streamable_http_mcp():
    """演示 Streamable HTTP MCP 的工作流程"""
    print(" Streamable HTTP MCP 客户端演示")
    print("=" * 50)

    async with StreamableHTTPMCPClient() as client:
        # 1. 初始化连接
        print("\n 1. 初始化连接...")
        if not await client.initialize():
            print(" 无法继续演示")
            return

        # 2. 获取工具列表
        print("\n 2. 获取工具列表...")
        tools = await client.list_tools()
        if not tools:
            print(" 无法获取工具列表")
            return

        # 3. 调用简单工具
        print("\n 3. 调用简单工具 (echo)...")
        result = await client.call_tool(
            "echo", {"message": "Hello, Streamable HTTP MCP!"}
        )
        if result:
            print(f"   结果: {result}")

        # 4. 调用长运行工具（非流式）
        print("\n 4. 调用长运行工具（非流式）...")
        result = await client.call_tool(
            "streaming_task", {"duration": 3, "stream": False}
        )
        if result:
            print(f"   结果: {result}")

        # 5. 调用长运行工具（流式）
        print("\n 5. 调用长运行工具（流式）...")
        result = await client.call_tool(
            "streaming_task", {"duration": 3, "stream": True}
        )
        if result:
            print(f"   结果: {result}")

        # 6. 尝试调用有状态工具
        print("\n 6. 尝试调用有状态工具...")
        result = await client.call_tool("stateful_counter", {"action": "increment"})
        if result:
            if "error" in result:
                print(f"   预期结果: {result['error']['message']}")
            else:
                print(f"   结果: {result}")

        # 7. 建立 SSE 连接（如果可用）
        print("\n7️⃣ 建立 SSE 连接...")
        if await client.establish_sse_connection():
            # 等待一下让连接稳定并接收一些消息
            print("   ⏳ 等待 SSE 消息...")
            await asyncio.sleep(3)  # 增加等待时间

            # 手动取消 SSE 任务，避免在 shutdown 时出现异常
            if client.sse_task:
                client.sse_task.cancel()
                try:
                    await client.sse_task
                except asyncio.CancelledError:
                    pass

        # 8. 关闭连接
        print("\n 8. 关闭连接...")
        await client.shutdown()

    print("\n" + "=" * 50)
    print(" Streamable HTTP 传输的特点总结:")
    print("   支持无状态和有状态模式")
    print("   统一的 /message 端点")
    print("   可升级的请求（支持 SSE）")
    print("   灵活的会话管理")
    print("   与现有 HTTP 基础设施完全兼容")
    print("   渐进式功能升级")
    print("   支持请求-响应模式")
    print("   可选的流式传输")


async def compare_transports():
    """比较两种传输方式的差异"""
    print("\n" + "=" * 60)
    print(" 传输方式对比分析")
    print("=" * 60)

    print("\n 传统 HTTP+SSE vs Streamable HTTP")
    print("-" * 40)

    print("\n 传统 HTTP+SSE 的局限性:")
    print("    连接断开后无法恢复")
    print("    服务器必须维护长连接")
    print("    对服务器基础设施要求高")
    print("    只能通过 SSE 发送服务器消息")
    print("    需要单独的 /sse 端点")

    print("\n Streamable HTTP 的优势:")
    print("   支持无状态服务器")
    print("   统一的 /message 端点")
    print("   可升级的请求（支持 SSE）")
    print("   灵活的会话管理")
    print("   与现有 HTTP 基础设施完全兼容")
    print("   渐进式功能升级")
    print("   支持请求-响应模式")

    print("\n 适用场景:")
    print("   传统 HTTP+SSE:")
    print("     - 需要实时双向通信")
    print("     - 服务器基础设施完善")
    print("     - 对连接稳定性要求高")

    print("\n   Streamable HTTP:")
    print("     - 需要灵活部署选项")
    print("     - 基础设施兼容性要求高")
    print("     - 需要渐进式功能升级")
    print("     - 支持无状态部署")


if __name__ == "__main__":
    print("启动 Streamable HTTP MCP 客户端演示...")
    print(" 请确保服务器已启动: python streamable_server.py")
    print()

    try:
        asyncio.run(demo_streamable_http_mcp())
        asyncio.run(compare_transports())
    except KeyboardInterrupt:
        print("\n 演示被用户中断")
    except Exception as e:
        print(f"\n 演示异常: {e}")
